# core/database_decorator.py
from functools import wraps
from sqlalchemy.ext.asyncio import AsyncSession
from typing import Callable, Any
import logging
from decorator import decorator

# 从你的数据库配置文件中导入已有的会话工厂
from ..dependencies import AsyncSessionLocal

logger = logging.getLogger(__name__)

def with_async_session(auto_commit: bool = True):
    """
    一个装饰器，用于自动管理异步SQLAlchemy会话的生命周期。
    
    参数:
        auto_commit (bool): 是否在执行成功后自动提交事务。默认为True。
                            设置为False时，需在函数内手动调用session.commit()。
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs) -> Any:
            # 创建一个新的数据库会话
            session: AsyncSession = AsyncSessionLocal()
            # 关键步骤：将session注入到被装饰函数的参数中
            kwargs['session'] = session
            
            try:
                # 执行被装饰的函数
                result = await func(*args, **kwargs)
                # 如果启用自动提交且会话中有事务，则提交
                if auto_commit and session.in_transaction():
                    await session.commit()
                return result
                
            except Exception as e:
                # 发生异常时回滚事务
                if session.in_transaction():
                    await session.rollback()
                logger.error(f"Database operation failed in {func.__name__}: {e}")
                raise e  # 重新抛出异常，让上层调用者处理
                
            finally:
                # 无论成功与否，最终都会确保会话被关闭
                await session.close()
                
        return wrapper
    return decorator